from typing import Dict, Any, List
import openai
from agents.components.embed.base_embedding_component import EmbeddingComponent
from agents.components.base import ComponentConfig

class OpenAIEmbeddingComponent(EmbeddingComponent):
    """使用 OpenAI 进行文本向量化的组件"""
    
    @classmethod
    def get_component_name(cls) -> str:
        return "openai_embedding"
    
    @classmethod
    def get_component_description(cls) -> str:
        return "使用 OpenAI 进行文本向量化的组件"
    
    @classmethod
    def get_config_schema(cls) -> Dict[str, Any]:
        base_schema = super().get_config_schema()
        base_schema["properties"].update({
            "api_key": {
                "type": "string",
                "description": "OpenAI API Key"
            },
            "model_name": {
                "type": "string",
                "description": "OpenAI 模型名称",
                "default": "text-embedding-3-small"
            },
            "organization": {
                "type": "string",
                "description": "OpenAI 组织 ID",
                "default": None
            }
        })
        return base_schema
    
    def __init__(self, config: ComponentConfig):
        super().__init__(config)
        self.api_key = config.config.get("api_key")
        self.model_name = config.config.get("model_name", "text-embedding-3-small")
        self.organization = config.config.get("organization")
        
        if not self.api_key:
            raise ValueError("OpenAI API Key 不能为空")
            
        openai.api_key = self.api_key
        if self.organization:
            openai.organization = self.organization
        
    async def _embed_texts(self, texts: List[str]) -> Dict[str, Any]:
        try:
            response = await openai.Embedding.acreate(
                model=self.model_name,
                input=texts
            )
            
            embeddings = [item["embedding"] for item in response.data]
            return {"embeddings": embeddings}
            
        except Exception as e:
            raise Exception(f"OpenAI API 调用失败: {str(e)}") 